import pandas as pd
import numpy as np
import os
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
import xgboost as xgb
import matplotlib.pyplot as plt
import seaborn as sns

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['figure.dpi'] = 300

# 定义数据路径和输出路径
DATA_PATH = r"D:\大三上\大数据分析及数据可视化\《Excel数据可视化 - 从图表到数据大屏》-清华-郭宏远\data-visualization\data\erp_order_data.xlsx"
OUTPUT_DIR = r"C:\Users\31417\Desktop"

# 确保输出目录存在
os.makedirs(OUTPUT_DIR, exist_ok=True)


def load_erp_data():
    """加载ERP订单数据，处理文件路径与日期格式"""
    try:
        df = pd.read_excel(DATA_PATH)
        # 转换日期列
        date_columns = ['shipping_date', 'payment_date', 'order_time']
        for col in date_columns:
            if col in df.columns:
                df[col] = pd.to_datetime(df[col])
        return df
    except Exception as e:
        print(f"数据加载失败: {e}")
        raise FileNotFoundError("未找到数据文件")


def prepare_order_level_data(df):
    """在订单级别准备数据，保留更多原始信息"""
    df = df.copy()

    # 1. 特征工程：从时间中提取特征
    df['year'] = df['order_time'].dt.year
    df['month'] = df['order_time'].dt.month
    df['dayofweek'] = df['order_time'].dt.dayofweek
    df['is_weekend'] = (df['dayofweek'] >= 5).astype(int)

    # 2. 选择用于独热编码的分类变量
    categorical_features = ['province', 'city', 'platform', 'store_name', 'product_name']

    # 3. 对分类变量进行独热编码
    df_encoded = pd.get_dummies(df, columns=categorical_features, prefix=categorical_features, drop_first=True)

    # 4. 选择最终用于建模的特征列（排除原始时间、ID、金额、状态等列）
    exclude_columns = [
        'id', 'internal_order_number', 'online_order_number',
        'full_channel_user_id', 'consignee', 'spu', 'sku',
        'color_and_spec', 'original_price', 'payable_amount',
        'paid_amount', 'shipping_date', 'payment_date',
        'order_time', 'product_amount',
        'status', 'sub_order_status', 'refund_status',
        'is_gift', 'sub_order_number', 'online_sub_order_number',
        'original_online_order_number', 'registered_quantity',
        'actual_refund_quantity'
    ]

    feature_columns = [col for col in df_encoded.columns if col not in exclude_columns]
    X = df_encoded[feature_columns]
    y = df_encoded['product_amount']

    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    return X_train, X_test, y_train, y_test


def train_random_forest_model(X_train, y_train):
    """训练随机森林模型"""
    model = RandomForestRegressor(
        n_estimators=200,
        max_depth=15,
        min_samples_split=10,
        random_state=42,
        n_jobs=-1
    )
    model.fit(X_train, y_train)
    return model


def train_xgboost_model(X_train, y_train):
    """训练XGBoost模型"""
    model = xgb.XGBRegressor(
        n_estimators=200,
        max_depth=8,
        learning_rate=0.1,
        subsample=0.8,
        colsample_bytree=0.8,
        random_state=42,
        n_jobs=-1
    )
    model.fit(X_train, y_train)
    return model


def evaluate_model(model, X_test, y_test, model_name, output_dir):
    """评估模型性能并保存多张评估图表"""
    y_pred = model.predict(X_test)
    mse = mean_squared_error(y_test, y_pred)
    rmse = np.sqrt(mse)
    r2 = r2_score(y_test, y_pred)

    print(f"{model_name} 模型评估结果:")
    print(f"RMSE: {rmse:.2f}")
    print(f"R²: {r2:.4f}")

    # === 图表1: 真实值 vs 预测值散点图 ===
    plt.figure(figsize=(10, 8))
    plt.scatter(y_test, y_pred, alpha=0.6, color='#4BB5C2', s=30)
    plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=2)
    plt.title(f'{model_name} 模型：真实值 vs 预测值', fontsize=14)
    plt.xlabel('真实销售额')
    plt.ylabel('预测销售额')
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f'{model_name}_01_true_vs_pred.png'), dpi=300, bbox_inches='tight')
    plt.show()

    # === 图表2: 残差图 (Residuals Plot) ===
    residuals = y_test - y_pred
    plt.figure(figsize=(10, 8))
    plt.scatter(y_pred, residuals, alpha=0.6, color='#E56A72', s=30)
    plt.axhline(y=0, color='r', linestyle='--')
    plt.title(f'{model_name} 模型：残差图', fontsize=14)
    plt.xlabel('预测销售额')
    plt.ylabel('残差 (真实值 - 预测值)')
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f'{model_name}_02_residuals.png'), dpi=300, bbox_inches='tight')
    plt.show()

    return rmse, r2


def feature_importance_analysis(model, X_train, model_name, output_dir):
    """分析特征重要性并保存图表"""
    if hasattr(model, 'feature_importances_'):
        importances = model.feature_importances_
        feature_names = X_train.columns

        # 创建DataFrame以便排序
        importance_df = pd.DataFrame({
            'feature': feature_names,
            'importance': importances
        }).sort_values('importance', ascending=False)

        # 只取前15个最重要的特征进行可视化
        top_features = importance_df.head(15)

        # === 图表3: 特征重要性条形图 ===
        plt.figure(figsize=(10, 8))
        plt.barh(range(len(top_features)), top_features['importance'], align='center')
        plt.yticks(range(len(top_features)), top_features['feature'])
        plt.gca().invert_yaxis()
        plt.title(f'{model_name} 模型特征重要性 (Top 15)', fontsize=14)
        plt.xlabel('重要性')
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, f'{model_name}_03_feature_importance.png'), dpi=300, bbox_inches='tight')
        plt.show()

        # 打印重要特征
        print(f"\n{model_name} 模型最重要的15个特征:")
        for idx, row in top_features.iterrows():
            print(f"{row['feature']}: {row['importance']:.4f}")

        return importance_df
    else:
        print(f"{model_name} 模型不支持特征重要性分析")
        return None


def sales_forecasting(df, output_dir):
    """销售预测主流程（订单级）"""
    print("开始销售预测分析（订单级别）...")

    # 准备数据
    X_train, X_test, y_train, y_test = prepare_order_level_data(df)
    print(f"训练集大小: {X_train.shape[0]}，测试集大小: {X_test.shape[0]}")
    print(f"总特征数量: {X_train.shape[1]}")

    # 训练随机森林模型
    print("\n训练随机森林模型...")
    rf_model = train_random_forest_model(X_train, y_train)
    rf_rmse, rf_r2 = evaluate_model(rf_model, X_test, y_test, "RandomForest", output_dir)
    rf_importance = feature_importance_analysis(rf_model, X_train, "RandomForest", output_dir)

    # 训练XGBoost模型
    print("\n训练XGBoost模型...")
    xgb_model = train_xgboost_model(X_train, y_train)
    xgb_rmse, xgb_r2 = evaluate_model(xgb_model, X_test, y_test, "XGBoost", output_dir)
    xgb_importance = feature_importance_analysis(xgb_model, X_train, "XGBoost", output_dir)

    # 选择最佳模型
    if xgb_r2 > rf_r2:
        best_model = xgb_model
        best_model_name = "XGBoost"
        best_importance = xgb_importance
        print(f"\nXGBoost模型表现更好，R² = {xgb_r2:.4f}")
    else:
        best_model = rf_model
        best_model_name = "RandomForest"
        best_importance = rf_importance
        print(f"\n随机森林模型表现更好，R² = {rf_r2:.4f}")

    # 保存完整的特征重要性到CSV
    if best_importance is not None:
        importance_path = os.path.join(output_dir, f'{best_model_name}_full_feature_importance.csv')
        best_importance.to_csv(importance_path, index=False, encoding='utf-8-sig')
        print(f"\n完整的特征重要性已保存至: {importance_path}")

    return best_model, best_model_name


if __name__ == "__main__":
    # 加载数据
    df = load_erp_data()
    print("数据加载完成！")

    # 执行销售预测
    best_model, best_model_name = sales_forecasting(df, OUTPUT_DIR)

    print("\n销售预测分析完成！")